"""工具间共享状态."""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Dict, Optional

import pandas as pd

from ..data.schemas import ETFHistoricalDataset


@dataclass
class DataStore:
    """简单的内存级数据缓存."""

    datasets: Dict[str, ETFHistoricalDataset] = field(default_factory=dict)
    feature_frames: Dict[str, pd.DataFrame] = field(default_factory=dict)
    risk_summaries: Dict[str, Dict[str, object]] = field(default_factory=dict)

    def set_dataset(self, symbol: str, dataset: ETFHistoricalDataset) -> None:
        """设置数据集缓存."""
        self.datasets[symbol] = dataset

    def get_dataset(self, symbol: str) -> Optional[ETFHistoricalDataset]:
        """获取数据集缓存."""
        return self.datasets.get(symbol)

    def set_features(self, symbol: str, features: pd.DataFrame) -> None:
        """设置特征数据缓存."""
        self.feature_frames[symbol] = features

    def get_features(self, symbol: str) -> Optional[pd.DataFrame]:
        """获取特征数据缓存."""
        return self.feature_frames.get(symbol)

    def set_risk_summary(self, symbol: str, summary: Dict[str, object]) -> None:
        """设置风险摘要缓存."""
        self.risk_summaries[symbol] = summary

    def get_risk_summary(self, symbol: str) -> Optional[Dict[str, object]]:
        """获取风险摘要缓存."""
        return self.risk_summaries.get(symbol)


GLOBAL_DATA_STORE = DataStore()

